import copy
from typing import List, Union

import torch
from torch import Tensor

from individual.solution import Solution
from utils.vectors import unique


class Population:
    def __init__(self, dec: Union[Tensor, list[Solution]] = None, obj: Tensor = None, **kwargs):
        self.populations = []
        if isinstance(dec, list):
            self.populations = dec
            return
        if dec is None or obj is None:
            return
        [pop_size, _] = dec.shape
        for i in range(pop_size):
            solution = Solution(dec[i], obj[i])
            for key, values in kwargs.items():
                solution.__setattr__(key, values[i])
            self.populations.append(solution)

    @property
    def dec(self):
        return self.get_attr('dec')

    @property
    def obj(self):
        return self.get_attr('obj')

    @property
    def con(self):
        return self.get_attr('con')

    # def __getattr__(self, item):
    #     if item == 'con':
    #         return torch.zeros(len(self.populations), 1, dtype=torch.double)
    #     return super().__getattr__(item)

    def __setitem__(self, index: Union[int, Tensor], value: Union[Solution, 'Population']):
        if isinstance(value, Solution) and isinstance(index, Tensor) and len(index) > 1:
            raise ValueError(" size index != len(value) ")
        if isinstance(index, slice):
            value = value.populations
            # 计算切片的长度
            start, stop, step = index.indices(len(self.populations))
            slice_length = max(0, (stop - start + (step - (1 if step > 0 else -1))) // step)
            # 检查值的长度是否与切片长度匹配
            if len(value) != slice_length:
                raise ValueError(f"Attempt to assign a sequence of size {len(value)} to a slice of size {slice_length}")
            # 处理切片赋值
            for i, item in enumerate(range(start, stop, step)):
                self.populations[item] = value[i]
        elif isinstance(index, Tensor):
            value = value.populations
            # 处理布尔tensor索引
            if index.dtype == torch.bool:
                if index.numel() != len(self.populations):
                    raise ValueError("The length of the boolean index does not match the length of the data list.")
                # 使用布尔索引找到对应的索引位置
                idx_list = [i for i, x in enumerate(index) if x]
            # 处理整数tensor索引
            elif index.dtype == torch.int64 or index.dtype == torch.int32:
                idx_list = index.tolist()
            else:
                raise TypeError("Index tensor must be either integer or boolean type.")

            # 应用索引列表对data进行操作
            for i, idx in enumerate(idx_list):
                self.populations[idx] = value[i]
        else:
            self.populations[index] = value

    def get_attr(self, attr_name: str) -> Tensor:
        """
        根据属性名获取种群的属性
        :param attr_name:
        :return: 返回所有个体的该属性组成的二维张量
        """
        try:
            return torch.cat(
                list(map(lambda solution: getattr(solution, attr_name).unsqueeze(0), self.populations)), dim=0)
        except AttributeError:
            raise AttributeError(f"there is no name for {attr_name} property")

    # 将满足索引的individual组成list赋值给self.populations
    def screen_populations(self, index: Tensor):
        self.populations = self[index].populations

    def __getitem__(self, index: Union[Tensor, int]) -> 'Population':
        if isinstance(index, int):
            return self.populations[index]
        elif isinstance(index, Tensor):
            index = index.numpy()
            populations = []
            if index.dtype == bool:
                if len(self.populations) != len(index):
                    raise ValueError('bool索引大小与种群大小不匹配')
                populations = [individual for individual, include in zip(self.populations, index) if include]
            else:
                for i in range(len(index)):
                    populations.append(self.populations[index[i]])
        else:
            raise TypeError("索引类型错误，只支持bool或者int类型的tensor索引")
        return Population(populations)

    # 重写自增操作
    def __iadd__(self, other: Union[Solution, 'Population']):
        if isinstance(other, Population):
            self.populations.extend(other.populations)
        elif isinstance(other, Solution):
            self.populations.extend(other)
        else:
            raise TypeError(f"不支持的类型 {type(other)} 进行iadd操作")
        return self

    def __add__(self, other: Union['Population', Solution]):
        population = Population()
        if isinstance(other, Population):
            population.populations = copy.deepcopy(self.populations) + other.populations
            index = unique(population.obj)
            return population[index]
        elif isinstance(other, Solution):
            population.populations = copy.deepcopy(self.populations)
            population.populations.append(other)
        else:
            raise TypeError(f"不支持的类型 {type(other)} 进行add操作")
        return population

    def __len__(self):
        return len(self.populations)

    def replace_individual(self, population: 'Population', index: Tensor):
        """
        替换种群中的部分解
        :param population:用于替换的种群
        :param index:替换的索引
        :return:
        """
        for index, value in enumerate(index):
            self.populations[value] = population.populations[index]

    def set_populations(self, populations: List[Solution]):
        self.populations = populations

    def get_best_solution(self):
        from algorithms.base.nsga2 import NSGA2
        """
        返回种群的非支配解组成的二维张量
        :return:
        """
        obj = self.get_attr("obj")
        front_no, _ = NSGA2.ND_sort(obj, 1)
        return obj[front_no == 1]
